%	Compute a case with 3 states of LANCIA, RUSSO, and WORRALL (2024)                  %
%   The value iteration is using ex-ante promised utility as state variable            %
%	Calls functions:    Utility.m                                                      %
%                       Autarky.m                                                      %
%                       Value_Iteration_3state.m                                       %

clear all
tic
options = optimset('Display','iter','LargeScale','off','MaxIter',1000,'TolFun',1e-10,'TolX',1e-10); 

cd ..\
cd Matlab_functions

%----------------------------------------------
% Parameter Setting
%----------------------------------------------

delta = exp(-1/75);                        % Discount Factor
gamma = 1;                            % Risk Aversion
pi3 = 0.25;                           % prob state 1
pi2 = 0.25;                           % prob state 2
pi1 = 1-pi3-pi2;                      % prob state 3
e1 = 1;                               % endowment state 1
e2 = 1;                               % endowment state 2
e3 = 1;                               % endowment state 3
sigma3 = 0.05*(2.5)^2;                % measure of volatility of endowment
sigma2 = 0.05*(2.5);                  % measure of volatility of endowment
sigma1 = 0;                           % measure of volatility of endowment
mu = 0.5;                             % measure of mean of endowment
ey3 = (mu+sigma3)*e3;                  % endowment state 1 young
eo3 = (1-mu-sigma3)*e3;               % endowment state 1 old
ey2 = (mu+sigma2)*e2;                  % endowment state 2 young
eo2 = (1-mu-sigma2)*e2;               % endowment state 2 old
ey1 = (mu+sigma1)*e1;                  % endowment state 3 young
eo1 = (1-mu-sigma1)*e1;               % endowment state 3 old
betta = delta;

%----------------------------------------------
% Checks
%----------------------------------------------
lamb1 = (eo1/ey1)^gamma;
lamb2 = (eo2/ey2)^gamma;
lamb3 = (eo3/ey3)^gamma;
S1 = lamb1-betta;          % this must be positive, state 1 bad state for the young
S2 = lamb2-betta;          % this can be positive or negative
S3 = lamb3-betta;          % this must be negative, state 3 good state for the young

Sfb = lamb1-betta/delta;   % if negative, then FB implies positive transfers in all states
Nec_cond = betta*(pi3*(ey3/eo3)^gamma+(pi2)*(ey2/eo2)^gamma+(pi1)*(ey1/eo1)^gamma)-1; % if positive, then dynamic efficiency in deterministic case

%----------------------------------------------
% Find max and min of promised values
%----------------------------------------------
v1aut = (utility(ey1,gamma)+betta*(pi1*utility(eo1,gamma)+(pi2)*utility(eo2,gamma)+(pi3)*utility(eo3,gamma))); % intertemporal utility in autarky in state 1
v2aut = (utility(ey2,gamma)+betta*(pi1*utility(eo1,gamma)+(pi2)*utility(eo2,gamma)+(pi3)*utility(eo3,gamma))); % intertemporal utility in autarky in state 2
v3aut = (utility(ey3,gamma)+betta*(pi1*utility(eo1,gamma)+(pi2)*utility(eo2,gamma)+(pi3)*utility(eo3,gamma))); % intertemporal utility in autarky in state 2
waut = pi1*utility(eo1,gamma)+(pi2)*utility(eo2,gamma)+(pi3)*utility(eo3,gamma);

st = @(z)[utility(ey1-z(1),gamma)+betta*(pi1*utility(eo1+z(1),gamma)+(pi2)*utility(eo2+z(2),gamma)+(pi3)*utility(eo3+z(3),gamma))-v1aut;      % IC state 1
        utility(ey2-z(2),gamma)+betta*(pi1*utility(eo1+z(1),gamma)+(pi2)*utility(eo2+z(2),gamma)+(pi3)*utility(eo3+z(3),gamma))-v2aut;
        utility(ey3-z(3),gamma)+betta*(pi1*utility(eo1+z(1),gamma)+(pi2)*utility(eo2+z(2),gamma)+(pi3)*utility(eo3+z(3),gamma))-v3aut
        ];       

Zst = fsolve(st,[0.2 0.2 0.2]);

tau1_max = Zst(1);      % consumption of old state 1
tau2_max = Zst(2);      % consumption of old state 2
tau3_max = Zst(3);      % consumption of old state 3
w_st = pi1*utility(eo1+tau1_max,gamma)+pi2*utility(eo2+tau2_max,gamma)+pi3*utility(eo3+tau3_max,gamma); 

%----------------------------------------------
% grid for promised values
%----------------------------------------------
wmin = waut;
wmax = w_st;
nw = 150;
wgrid = cheby_points(wmin,wmax,nw);

param = [delta,betta,gamma,pi1,pi2,pi3,e1,e2,e3,ey1,ey2,ey3,v1aut,v2aut,v3aut,wmin,wmax,nw];
%----------------------------------------------
% Iterating the value Function
%----------------------------------------------
P_old=0.1*ones(nw,1);
distmia=100;
olddist=1;
iter=0;
tollv=10^(-4);
tic
fprintf('For this tolerance, I predict that about %g iterations will be needed \n',round(fsolve(@(x)tollv.^(1/x)-delta,3)));

cy1_old = (ey1-0.001)*ones(nw,1);
cy2_old = (ey2-0.001)*ones(nw,1);
cy3_old = (ey3-0.001)*ones(nw,1);
w1_old = 0.1*ones(nw,1);
w2_old = 0.1*ones(nw,1);
w3_old = 0.1*ones(nw,1);

% value iteration 
while distmia > tollv
    iter=iter+1;
    
    [P_new,cy1,cy2,cy3,w1,w2,w3] = Value_Iteration_3state(wgrid,P_old,cy1_old,cy2_old,cy3_old,w1_old,w2_old,w3_old,param);
    distmia=max(abs(P_new-P_old));
    
    fprintf('Iteration %g completed , Distance is: %g   Improved by %9.1f %% \n',iter,distmia,(((olddist-distmia)/olddist))*100);
    
    figure(1)
    subplot(2,2,1)
             plot(wgrid,w1,'b'); hold on;...
             plot(wgrid,w2,'--r'); hold on;...
             plot(wgrid,w3,'--g'); hold on;...
             plot(wgrid,wgrid,'black'); hold on;
                    title('Promised utilities','fontsize',08)
    subplot(2,2,2)                
             plot(wgrid,P_new,'b'); hold on;
                    title('Value','fontsize',08)
                    
    wgridd=wgrid(2:length(wgrid));
    DP_1=diff(P_new')./diff(wgrid);
    subplot(2,2,3)
             plot(wgridd,DP_1); hold off; 
                   title('Derivative Value','fontsize',08)
    
    drawnow;
    P_old=P_new;
    olddist=distmia;
    cy1_old=cy1;
    cy2_old=cy2;
    cy3_old=cy3;
    w1_old=w1;
    w2_old=w2;
    w3_old=w3;
end

% Ex-post promise
% next period consumption
cy11=interp1(wgrid,cy1,w1,'spline');
cy12=interp1(wgrid,cy2,w1,'spline');
cy13=interp1(wgrid,cy3,w1,'spline');
cy21=interp1(wgrid,cy1,w2,'spline');
cy22=interp1(wgrid,cy2,w2,'spline');
cy23=interp1(wgrid,cy3,w2,'spline');
cy31=interp1(wgrid,cy1,w3,'spline');
cy32=interp1(wgrid,cy2,w3,'spline');
cy33=interp1(wgrid,cy3,w3,'spline');

omega1=utility(1-cy1,gamma);
omega2=utility(1-cy2,gamma);
omega3=utility(1-cy3,gamma);

omega11=utility(1-cy11,gamma);
omega12=utility(1-cy12,gamma);
omega13=utility(1-cy13,gamma);

omega21=utility(1-cy21,gamma);
omega22=utility(1-cy22,gamma);
omega23=utility(1-cy23,gamma);

omega31=utility(1-cy31,gamma);
omega32=utility(1-cy32,gamma);
omega33=utility(1-cy33,gamma);

cy1omega=interp1(wgrid,cy1,omega1,'spline');
cy2omega=interp1(wgrid,cy2,omega2,'spline');
cy3omega=interp1(wgrid,cy3,omega3,'spline');


%----------------------------------------------
% Outcome Variables
%----------------------------------------------

cd ..\
cd Stored_file
save('Policy_3state.mat', 'wgrid', 'P_new', 'w1', 'w2', 'w3', 'cy1', 'cy2', 'cy3', 'lamb1', 'lamb2', 'lamb3', 'delta', 'betta', 'gamma', 'pi1', 'pi2', 'pi3', 'e1', 'e2', 'e3', 'ey1', 'ey2', 'ey3', 'param')
